package org.example.util;

import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiModerationModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import org.example.config.OpenAiConfig;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

@Component
public class CreateModel {

    @Autowired
    OpenAiConfig openAiConfig;

    public StreamingChatLanguageModel createStreamingChatLanguageModel(){
        return OpenAiStreamingChatModel.builder()
                .baseUrl(openAiConfig.getBaseUrl())
                .apiKey(openAiConfig.getApiKey())
                .modelName("gpt-4")
                .temperature(0.1)
                .build();
    }

    public ChatLanguageModel createChatLanguageModel(){
        return OpenAiChatModel.builder()
                .baseUrl(openAiConfig.getBaseUrl())
                .apiKey(openAiConfig.getApiKey())
                .modelName("gpt-4")
                .temperature(0.1)
                .build();
    }

    public OpenAiEmbeddingModel createEmbeddingModel(){
        return OpenAiEmbeddingModel.builder()
                .baseUrl(openAiConfig.getBaseUrl())
                .apiKey(openAiConfig.getApiKey())
                .build();
    }

    public OpenAiModerationModel createModerationModel(){
        return OpenAiModerationModel.builder()
                .modelName("omni-moderation-latest")
                .apiKey(openAiConfig.getApiKey())
                .baseUrl(openAiConfig.getBaseUrl())
                .build();
    }
}
